from SentModel.Sent2Vec import *

tr = LMReader("../data/twitter_LM_tr.txt")
dev = LMReader("../data/twitter_LM_dev.txt")
te = LMReader("../data/twitter_LM_te.txt")
tr_loader = DataLoader(tr, batch_size=32, shuffle=True)
dev_loader = DataLoader(dev, batch_size=100, shuffle=True)
te_loader = DataLoader(te, batch_size=100, shuffle=True)

model = BertVec("../../bert_en/")
print("<========Distill Train=======>")
s_model = LSTMVec("../saved/bert_embedding.pkl", "../../bert_en/")
trainer = Distill_LM_Trainer(model, s_model, model.bert.config)
trainer.DistillLM_Train(tr_loader, dev_loader, te_loader, print_every=10, model_file="../saved/distill_LSTM_LM.pkl")